#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识库搜索模块
提供多种搜索方式：实体搜索、关系搜索、语义搜索等
"""

import json
import re
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class SearchResult:
    """搜索结果"""
    id: str
    type: str  # entity, relation, chunk
    title: str
    content: str
    score: float
    metadata: Dict[str, Any]

class KnowledgeBaseSearch:
    """知识库搜索引擎"""
    
    def __init__(self, knowledge_base_path: str = "knowledge_base.json"):
        """
        初始化搜索引擎
        
        Args:
            knowledge_base_path: 知识库文件路径
        """
        self.knowledge_base_path = knowledge_base_path
        self.knowledge_base = None
        self.load_knowledge_base()
    
    def load_knowledge_base(self) -> None:
        """加载知识库"""
        try:
            with open(self.knowledge_base_path, 'r', encoding='utf-8') as f:
                self.knowledge_base = json.load(f)
            logger.info(f"知识库已加载: {self.knowledge_base_path}")
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")
            raise
    
    def search_entities(self, query: str, entity_type: Optional[str] = None) -> List[SearchResult]:
        """
        搜索实体
        
        Args:
            query: 搜索查询
            entity_type: 实体类型过滤
            
        Returns:
            List[SearchResult]: 搜索结果列表
        """
        logger.info(f"搜索实体: {query}")
        
        results = []
        query_lower = query.lower()
        
        for entity_id, entity in self.knowledge_base["entities"].items():
            # 类型过滤
            if entity_type and entity["type"] != entity_type:
                continue
            
            # 计算匹配分数
            score = 0.0
            
            # 名称匹配
            if query_lower in entity["name"].lower():
                score += 0.8
            
            # 描述匹配
            if query_lower in entity["description"].lower():
                score += 0.6
            
            # 关键词匹配
            for keyword in entity["keywords"]:
                if query_lower in keyword.lower():
                    score += 0.4
            
            if score > 0:
                result = SearchResult(
                    id=entity_id,
                    type="entity",
                    title=entity["name"],
                    content=entity["description"],
                    score=score,
                    metadata={
                        "entity_type": entity["type"],
                        "keywords": entity["keywords"],
                        "relations": entity["relations"],
                        "confidence": entity["confidence"]
                    }
                )
                results.append(result)
        
        # 按分数排序
        results.sort(key=lambda x: x.score, reverse=True)
        logger.info(f"实体搜索完成，找到 {len(results)} 个结果")
        return results
    
    def search_relations(self, query: str, predicate: Optional[str] = None) -> List[SearchResult]:
        """
        搜索关系
        
        Args:
            query: 搜索查询
            predicate: 关系谓词过滤
            
        Returns:
            List[SearchResult]: 搜索结果列表
        """
        logger.info(f"搜索关系: {query}")
        
        results = []
        query_lower = query.lower()
        
        for relation_id, relation in self.knowledge_base["relations"].items():
            # 谓词过滤
            if predicate and relation["predicate"] != predicate:
                continue
            
            # 计算匹配分数
            score = 0.0
            
            # 主体匹配
            if query_lower in relation["subject"].lower():
                score += 0.5
            
            # 客体匹配
            if query_lower in relation["object"].lower():
                score += 0.5
            
            # 谓词匹配
            if query_lower in relation["predicate"].lower():
                score += 0.8
            
            # 描述匹配
            if query_lower in relation["description"].lower():
                score += 0.6
            
            if score > 0:
                result = SearchResult(
                    id=relation_id,
                    type="relation",
                    title=f"{relation['subject']} {relation['predicate']} {relation['object']}",
                    content=relation["description"],
                    score=score,
                    metadata={
                        "subject": relation["subject"],
                        "predicate": relation["predicate"],
                        "object": relation["object"],
                        "confidence": relation["confidence"]
                    }
                )
                results.append(result)
        
        # 按分数排序
        results.sort(key=lambda x: x.score, reverse=True)
        logger.info(f"关系搜索完成，找到 {len(results)} 个结果")
        return results
    
    def search_chunks(self, query: str, top_k: int = 10) -> List[SearchResult]:
        """
        搜索知识块
        
        Args:
            query: 搜索查询
            top_k: 返回结果数量
            
        Returns:
            List[SearchResult]: 搜索结果列表
        """
        logger.info(f"搜索知识块: {query}")
        
        results = []
        query_lower = query.lower()
        
        for chunk_id, chunk in self.knowledge_base["chunks"].items():
            # 计算匹配分数
            score = 0.0
            
            # 标题匹配
            if query_lower in chunk["title"].lower():
                score += 0.8
            
            # 内容匹配
            content_lower = chunk["content"].lower()
            if query_lower in content_lower:
                score += 0.6
            
            # 关键词匹配
            for keyword in chunk["keywords"]:
                if query_lower in keyword.lower():
                    score += 0.4
            
            # 实体匹配
            for entity_id in chunk["entities"]:
                if entity_id in self.knowledge_base["entities"]:
                    entity = self.knowledge_base["entities"][entity_id]
                    if query_lower in entity["name"].lower():
                        score += 0.3
            
            if score > 0:
                result = SearchResult(
                    id=chunk_id,
                    type="chunk",
                    title=chunk["title"],
                    content=chunk["content"],
                    score=score,
                    metadata={
                        "keywords": chunk["keywords"],
                        "entities": chunk["entities"],
                        "relations": chunk["relations"],
                        "page_number": chunk["page_number"],
                        "chunk_index": chunk["chunk_index"]
                    }
                )
                results.append(result)
        
        # 按分数排序并限制数量
        results.sort(key=lambda x: x.score, reverse=True)
        results = results[:top_k]
        
        logger.info(f"知识块搜索完成，找到 {len(results)} 个结果")
        return results
    
    def search_all(self, query: str, top_k: int = 10) -> Dict[str, List[SearchResult]]:
        """
        全库搜索
        
        Args:
            query: 搜索查询
            top_k: 每种类型返回结果数量
            
        Returns:
            Dict[str, List[SearchResult]]: 分类搜索结果
        """
        logger.info(f"全库搜索: {query}")
        
        results = {
            "entities": self.search_entities(query)[:top_k],
            "relations": self.search_relations(query)[:top_k],
            "chunks": self.search_chunks(query, top_k)
        }
        
        total_results = sum(len(result_list) for result_list in results.values())
        logger.info(f"全库搜索完成，共找到 {total_results} 个结果")
        
        return results
    
    def get_entity_relations(self, entity_name: str) -> List[Dict[str, Any]]:
        """
        获取实体的所有关系
        
        Args:
            entity_name: 实体名称
            
        Returns:
            List[Dict[str, Any]]: 关系列表
        """
        logger.info(f"获取实体关系: {entity_name}")
        
        relations = []
        
        for relation_id, relation in self.knowledge_base["relations"].items():
            if (relation["subject"] == entity_name or 
                relation["object"] == entity_name):
                relations.append({
                    "id": relation_id,
                    "subject": relation["subject"],
                    "predicate": relation["predicate"],
                    "object": relation["object"],
                    "description": relation["description"],
                    "confidence": relation["confidence"]
                })
        
        logger.info(f"找到 {len(relations)} 个关系")
        return relations
    
    def get_related_entities(self, entity_name: str) -> List[Dict[str, Any]]:
        """
        获取相关实体
        
        Args:
            entity_name: 实体名称
            
        Returns:
            List[Dict[str, Any]]: 相关实体列表
        """
        logger.info(f"获取相关实体: {entity_name}")
        
        related_entities = set()
        
        # 通过关系找到相关实体
        for relation in self.knowledge_base["relations"].values():
            if relation["subject"] == entity_name:
                related_entities.add(relation["object"])
            elif relation["object"] == entity_name:
                related_entities.add(relation["subject"])
        
        # 获取实体详细信息
        result = []
        for entity_id, entity in self.knowledge_base["entities"].items():
            if entity["name"] in related_entities:
                result.append({
                    "id": entity_id,
                    "name": entity["name"],
                    "type": entity["type"],
                    "description": entity["description"],
                    "keywords": entity["keywords"]
                })
        
        logger.info(f"找到 {len(result)} 个相关实体")
        return result
    
    def get_knowledge_graph(self, center_entity: str, depth: int = 2) -> Dict[str, Any]:
        """
        获取知识图谱
        
        Args:
            center_entity: 中心实体
            depth: 图谱深度
            
        Returns:
            Dict[str, Any]: 知识图谱数据
        """
        logger.info(f"构建知识图谱: {center_entity}, 深度: {depth}")
        
        graph = {
            "nodes": [],
            "edges": [],
            "center_entity": center_entity
        }
        
        # 添加中心节点
        center_entity_info = None
        for entity_id, entity in self.knowledge_base["entities"].items():
            if entity["name"] == center_entity:
                center_entity_info = {
                    "id": entity_id,
                    "name": entity["name"],
                    "type": entity["type"],
                    "level": 0
                }
                graph["nodes"].append(center_entity_info)
                break
        
        if not center_entity_info:
            logger.warning(f"未找到中心实体: {center_entity}")
            return graph
        
        # 递归构建图谱
        self._build_graph_recursive(graph, center_entity, depth, 1, set())
        
        logger.info(f"知识图谱构建完成，节点数: {len(graph['nodes'])}, 边数: {len(graph['edges'])}")
        return graph
    
    def _build_graph_recursive(self, graph: Dict[str, Any], entity_name: str, 
                             max_depth: int, current_depth: int, visited: set) -> None:
        """递归构建知识图谱"""
        if current_depth > max_depth or entity_name in visited:
            return
        
        visited.add(entity_name)
        
        # 找到相关关系
        for relation_id, relation in self.knowledge_base["relations"].items():
            if relation["subject"] == entity_name:
                target_entity = relation["object"]
            elif relation["object"] == entity_name:
                target_entity = relation["subject"]
            else:
                continue
            
            # 添加目标节点
            target_entity_info = None
            for entity_id, entity in self.knowledge_base["entities"].items():
                if entity["name"] == target_entity:
                    target_entity_info = {
                        "id": entity_id,
                        "name": entity["name"],
                        "type": entity["type"],
                        "level": current_depth
                    }
                    break
            
            if target_entity_info and target_entity_info not in graph["nodes"]:
                graph["nodes"].append(target_entity_info)
            
            # 添加边
            edge = {
                "id": relation_id,
                "source": entity_name,
                "target": target_entity,
                "predicate": relation["predicate"],
                "description": relation["description"]
            }
            if edge not in graph["edges"]:
                graph["edges"].append(edge)
            
            # 递归处理
            self._build_graph_recursive(graph, target_entity, max_depth, 
                                      current_depth + 1, visited.copy())
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取知识库统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        stats = {
            "total_entities": len(self.knowledge_base["entities"]),
            "total_relations": len(self.knowledge_base["relations"]),
            "total_chunks": len(self.knowledge_base["chunks"]),
            "entity_types": {},
            "relation_types": {},
            "top_keywords": {}
        }
        
        # 统计实体类型
        for entity in self.knowledge_base["entities"].values():
            entity_type = entity["type"]
            stats["entity_types"][entity_type] = stats["entity_types"].get(entity_type, 0) + 1
        
        # 统计关系类型
        for relation in self.knowledge_base["relations"].values():
            predicate = relation["predicate"]
            stats["relation_types"][predicate] = stats["relation_types"].get(predicate, 0) + 1
        
        # 统计关键词
        keyword_freq = {}
        for chunk in self.knowledge_base["chunks"].values():
            for keyword in chunk["keywords"]:
                keyword_freq[keyword] = keyword_freq.get(keyword, 0) + 1
        
        stats["top_keywords"] = dict(sorted(keyword_freq.items(), 
                                          key=lambda x: x[1], reverse=True)[:10])
        
        return stats


def main():
    """测试知识库搜索功能"""
    search = KnowledgeBaseSearch()
    
    # 测试搜索
    test_queries = [
        "人与自然",
        "生态危机",
        "可持续发展",
        "工业革命"
    ]
    
    for query in test_queries:
        print(f"\n搜索查询: {query}")
        print("=" * 50)
        
        # 全库搜索
        results = search.search_all(query, top_k=3)
        
        for result_type, result_list in results.items():
            if result_list:
                print(f"\n{result_type.upper()} 结果:")
                for i, result in enumerate(result_list, 1):
                    print(f"  {i}. {result.title} (分数: {result.score:.2f})")
                    print(f"     {result.content[:100]}...")
    
    # 测试知识图谱
    print(f"\n知识图谱: 人与自然")
    print("=" * 50)
    graph = search.get_knowledge_graph("人与自然", depth=2)
    print(f"节点数: {len(graph['nodes'])}")
    print(f"边数: {len(graph['edges'])}")
    
    # 显示统计信息
    print(f"\n知识库统计:")
    print("=" * 50)
    stats = search.get_statistics()
    print(f"总实体数: {stats['total_entities']}")
    print(f"总关系数: {stats['total_relations']}")
    print(f"总知识块数: {stats['total_chunks']}")
    
    print(f"\n实体类型分布:")
    for entity_type, count in stats['entity_types'].items():
        print(f"  {entity_type}: {count}")
    
    print(f"\n关系类型分布:")
    for relation_type, count in stats['relation_types'].items():
        print(f"  {relation_type}: {count}")


if __name__ == "__main__":
    main()
